# %%
import datetime

import numpy as np
from matplotlib import pyplot as plt

from common.logger import TimeIt
from common.rs_utils.derived import compute_ndxi

from catalog.models.sentinel2_l2a import Sentinel2L2AScene
from catalog.raster import retrieve_by_item

# %%

items = Sentinel2L2AScene.query_many_items(
    temporal=(datetime.datetime(2021, 1, 1), datetime.datetime(2021, 1, 10)),
    mgrs_tile="48RVV",
)

item = list(items)[0]
print(item.uid, item.cloud_cover, item.start_datetime, item.created)

# %%

with TimeIt("retrieve data"):
    rds = retrieve_by_item(item, ["B08", "B04"])
print(rds.meta)

# %%
with TimeIt("deduce mask"):
    nda_mask = (rds.data[0, :, :] > 0) & (rds.data[1, :, :] > 0)

with TimeIt("calc ndvi, use numexpr"):
    nda_ndvi, nda_mask1 = compute_ndxi(
        rds.data[0, :, :],
        rds.data[1, :, :],
        nda_mask,
        # NOTE all the below are default args, just show here to provide more context for readers
        working_dtype=np.float32,
        clipped=True,
        use_ne=True,
    )

print(nda_ndvi.min(), nda_ndvi.max())

# %%

with TimeIt("calc ndvi, numpy only"):
    nda_ndvi, nda_mask2 = compute_ndxi(
        rds.data[0, :, :], rds.data[1, :, :], nda_mask, use_ne=False
    )
# %%

plt.imshow(nda_ndvi, vmin=-1, vmax=1, cmap="RdYlGn")

# %%
